﻿using Ndns.Packet.Abstractions;
using Ndns.Packet.Abstractions.ResourceRecords;
using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Net;
using System.Text;

namespace Ndns.Packet.Parser
{
    class DnsResourceRecordParser
    {
        private readonly DnsHostNameParser _hostNameParser;
        private readonly ArrayPool<byte> _arrayPool;

        public DnsResourceRecordParser(DnsHostNameParser hostNameParser) : this(hostNameParser, null) { }
        public DnsResourceRecordParser(DnsHostNameParser hostNameParser, ArrayPool<byte> arrayPool)
        {
            _hostNameParser = hostNameParser;
            _arrayPool = arrayPool;
        }

        public bool TryParse(Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset, out DnsResourceRecord result)
        {
            if (!_hostNameParser.TryParse(stream, symbols, ref offset, out var hostName))
            {
                result = null;
                return false;
            }

            var vType = stream.ReadNetworkUInt16();
            var vClass = stream.ReadNetworkUInt16();
            var vTTL = stream.ReadNetworkUInt32();
            if (!PacketEndOfStreamException.CheckValue(vTTL))
            {
                result = null;
                return false;
            }
            offset += 8;
            
            switch ((DnsRecordType)vType)
            {
                case DnsRecordType.A:
                    var recordA = new DnsResourceRecordA(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordA, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordA;
                    break;
                case DnsRecordType.NS:
                    var recordNS = new DnsResourceRecordNS(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordNS, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordNS;
                    break;
                case DnsRecordType.CNAME:
                    var recordCNAME = new DnsResourceRecordCNAME(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordCNAME, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordCNAME;
                    break;
                case DnsRecordType.SOA:
                    var recordSOA = new DnsResourceRecordSOA(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordSOA, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordSOA;
                    break;
                case DnsRecordType.PTR:
                    var recordPTR = new DnsResourceRecordPTR(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordPTR, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordPTR;
                    break;
                case DnsRecordType.MX:
                    var recordMX = new DnsResourceRecordMX(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordMX, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordMX;
                    break;
                case DnsRecordType.TXT:
                    var recordTXT = new DnsResourceRecordTXT(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordTXT, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordTXT;
                    break;
                case DnsRecordType.AAAA:
                    var recordAAAA = new DnsResourceRecordAAAA(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordAAAA, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordAAAA;
                    break;
                case DnsRecordType.DNAME:
                    var recordDNAME = new DnsResourceRecordDNAME(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordDNAME, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordDNAME;
                    break;
                case DnsRecordType.ANY:
                    var recordANY = new DnsResourceRecordAAAA(hostName, (DnsClass)vClass, (uint)vTTL);
                    if (!TryFillDataFrom(recordANY, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordANY;
                    break;
                default:
                    var recordCustom = new CustomDnsResourceRecord(hostName, (DnsClass)vClass, (DnsRecordType)vType, (uint)vTTL);
                    if (!TryFillDataFrom(recordCustom, stream, symbols, ref offset))
                    {
                        result = null;
                        return false;
                    }
                    result = recordCustom;
                    break;
            }
            return true;
        }


        public bool TryFillDataFrom(DnsResourceRecordA record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            offset += 2;
            if (len != 4)
                return false;
            using (var bufferLease = _arrayPool.LeaseOrAllocate(4))
            {
                var array = bufferLease.Array;
                Array.Clear(array, 0, 4);
                stream.Read(array, 0, 4);
                offset += 4;
                record.Address = new IPAddress(array);
            }
            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordAAAA record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            offset += 2;
            if (len != 16)
                return false;
            using (var bufferLease = _arrayPool.LeaseOrAllocate(16))
            {
                var array = bufferLease.Array;
                Array.Clear(array, 0, 16);
                stream.Read(array, 0, 16);
                offset += 16;
                record.Address = new IPAddress(array);
            }
            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordANY record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;
            var data = new byte[len];
            int readCount = stream.Read(data, 0, len);
            if (readCount != len)
                return false;
            offset += len;
            record.Data = data;
            return true;
        }

        public bool TryFillDataFrom(CustomDnsResourceRecord record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;
            var data = new byte[len];
            int readCount = stream.Read(data, 0, len);
            if (readCount != len)
                return false;
            offset += len;
            record.Data = data;
            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordCNAME record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;

            int originalOffset = offset;

            if (!_hostNameParser.TryParse(stream, symbols, ref offset, out var hostName))
                return false;
            record.CanonicalName = hostName;

            if (offset - originalOffset != len)
                return false;

            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordDNAME record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;

            int originalOffset = offset;

            if (!_hostNameParser.TryParse(stream, symbols, ref offset, out var hostName))
                return false;
            record.DelegationName = hostName;

            if (offset - originalOffset != len)
                return false;

            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordMX record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;

            int originalOffset = offset;
            
            var vPreference = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(vPreference))
                return false;
            offset += 2;
            record.Preference = (ushort)vPreference;

            if (!_hostNameParser.TryParse(stream, symbols, ref offset, out var hostName))
                return false;
            record.MXHostName = hostName;

            if (offset - originalOffset != len)
                return false;

            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordNS record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;

            int originalOffset = offset;

            if (!_hostNameParser.TryParse(stream, symbols, ref offset, out var hostName))
                return false;
            record.NameServerHostName = hostName;

            if (offset - originalOffset != len)
                return false;

            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordPTR record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;

            int originalOffset = offset;

            if (!_hostNameParser.TryParse(stream, symbols, ref offset, out var hostName))
                return false;
            record.ReversedDomainName = hostName;

            if (offset - originalOffset != len)
                return false;

            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordSOA record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;

            int originalOffset = offset;

            DnsHostName hostName;
            if (!_hostNameParser.TryParse(stream, symbols, ref offset, out hostName))
                return false;
            record.PrimaryNameServer = hostName;

            if (!_hostNameParser.TryParse(stream, symbols, ref offset, out hostName))
                return false;
            record.MailBox = hostName;

            var vSerialNumber = stream.ReadNetworkUInt32();
            var vRefreshInterval = stream.ReadNetworkUInt32();
            var vRetrialInterval = stream.ReadNetworkUInt32();
            var vExpireLimit = stream.ReadNetworkUInt32();
            var vMinimumTTL = stream.ReadNetworkUInt32();
            if (!PacketEndOfStreamException.CheckValue(vMinimumTTL))
                return false;
            offset += 20;

            unchecked
            {
                record.SerialNumber = (uint)vSerialNumber;
                record.RefreshInterval = (uint)vRefreshInterval;
                record.RetrialInterval = (uint)vRetrialInterval;
                record.ExpireLimit = (uint)vExpireLimit;
                record.MinimumTTL = (uint)vMinimumTTL;
            }
            
            if (offset - originalOffset != len)
                return false;

            return true;
        }

        public bool TryFillDataFrom(DnsResourceRecordTXT record, Stream stream, IDictionary<DnsHostNamePointer, DnsHostName> symbols, ref int offset)
        {
            Debug.Assert(record != null);

            int len = stream.ReadNetworkUInt16();
            if (!PacketEndOfStreamException.CheckValue(len))
                return false;
            offset += 2;

            if (len == 0)
            {
                record.Content = string.Empty;
                return true;
            }

            using (var bufferLease = _arrayPool.LeaseOrAllocate(len))
            {
                var buffer = bufferLease.Array;
                int copyCount = stream.CopyBytesTo(buffer, len);
                if (copyCount != len)
                    return false;
                offset += len;
                
                int bLen = buffer[0];
                if (bLen < 0 || bLen + 1 > len)
                    return false;
                record.Content = PunycodeHelper.DecodeDomainPart(Encoding.ASCII.GetString(buffer, 1, bLen));
            }

            return true;
        }
    }
}
